# pre packages
from myglobal import os, sys, LINES

# sys packages
from pylab import np
import obspy as ob
from glob import glob
from tqdm import tqdm
from obspy.core.utcdatetime import UTCDateTime
import pandas as pd
from scipy.stats import linregress


# self packages
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import norm, my_vel_regress
from utils.h5data import get_event_data, save_h5_data, read_h5_data, h5glob
from utils.hsr import RawDataBase, search_events
from utils.plot import plot_traces_by_subfigures, plot_events
from utils.trace import get_tapered_slices

# cmd
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-debug', action='store_true', help='method: DEBUG')
parser.add_argument('-base', default='DZ155', help='BASE_STAT, default  is DZ155')
parser.add_argument('-line', default='VFHF', help='LINE_ID, default is VFHF')
parser.add_argument('-S2N', action='store_true', help='direction of trains, default is S2N')
parser.add_argument('-N2S', action='store_true', help='direction of trains, default is N2S')
parser.add_argument('-emarker', default='', help='marker, ')
parser.add_argument('-amp', default=2e-4/4, type=float, help='amplitude threshold of hsr events')
parser.add_argument('-date', default='2303',  help='which data to use, 2303, 2406, 2409, 254C')
parser.add_argument('-input', default='',  help='input h5file')
parser.add_argument('-output', default='',  help='output csv file')

args = parser.parse_args()
print(args)

DEBUG= args.debug
LINE_ID = args.line
BASE_STAT=args.base
MARKER = args.emarker
S2N = args.S2N
N2S = args.N2S
DATE = args.date
HSR_EVENTT_AMP = args.amp

DATA_FILE = args.input
OUT_FILE = args.output


LINES = LINES[DATE]
if LINE_ID in LINES.keys():
    HSR_STATS = LINES[LINE_ID]
else:
    HSR_STATS = [LINE_ID]

if MARKER == '':
    MARKER=f'{BASE_STAT}.Line{LINE_ID}'
    if S2N:
        MARKER=MARKER+'S2N'
    if N2S:
        MARKER=MARKER+'N2S'
T_win =[-90,90]        
if DATE == '2303':
    H5_ROOT='data/2303.HSR/events'
    s_info = load_loc('./loc/loc_all_from_log_2303.csv')
    if N2S:
        T_win =[-40,60]
    if S2N:
        T_win =[-60,40]
elif DATE == '2406':
    pass
elif DATE == '2409':
    pass
elif DATE == '254C':
    H5_ROOT='data/254C.HSR/events'
    s_info = load_loc('./loc/254C.csv',lat_key='lat',lon_key='lon')
    if N2S:
        T_win =[-40,60]
    if S2N:
        T_win =[-60,40]
else:
    pass

FIG_ROOT = './figures/3.event.selected'
FIG_ROOT = f'{FIG_ROOT}/{MARKER}'
if not os.path.exists(FIG_ROOT):
    os.mkdir(FIG_ROOT)

x = []
for name in HSR_STATS:
    xi, _,_ = get_distance(s_info=s_info,name1=BASE_STAT,name2=name, S2N=True)
    x.append(xi)

x,HSR_STATS = sort_data_by_distance(x,NAMES=HSR_STATS)
x = np.array(x)
nx = len(x)

# groups = h5glob(DATA_FILE,pattern='DAY*/IDX*.T*', object_type='group')
# save_h5_data(DATA_FILE,{'all_groups':np.array(groups,dtype='S')}, mode='a')
# assert 1==2

groups = read_h5_data(DATA_FILE,keys=['all_groups'])[0]
groups = [g.decode() for g in groups]
groups.sort()
print(f'find {len(groups)} groups in {DATA_FILE}')
ne = len(groups)

IT = tqdm(range(ne),desc="select event data",total=ne)

event_select = pd.DataFrame(columns=["group","e_time","v1", "b1", "R1","v2", "b2", "R2"])
for i in IT:
    idx,group_i = i,groups[i]

    data_i, t, e_time, stats= read_h5_data(DATA_FILE, keys=['data','t','te','stats'],group_name=group_i)
    stats = [i.decode() for i in stats]
    e_time = e_time.decode()

    data_i,t = get_tapered_slices(data_i, t, T_win=T_win, L_Taper=1)
    
    t_measured =[]
    traces = []
    for j,name in enumerate(HSR_STATS):
        idx_j = stats.index(name)
        trace_j = data_i[idx_j,:]
        trace_j = np.abs(trace_j)
        
        amp_idx = np.where(trace_j>HSR_EVENTT_AMP)[0]
        if len(amp_idx)==0 or trace_j.max()==0:
            continue
        sp = min(amp_idx)
        ep = max(amp_idx)
        t_measured.append([x[j], t[sp],t[ep]])
        traces.append(data_i[idx_j,:])
    # print(t_measured)
    if len(t_measured)<3:
        continue
    if np.array(t_measured)[:,1].var()==0 or np.array(t_measured)[:,2].var()==0:
        continue

    t_measured = np.array(t_measured)
    x_valid = t_measured[:,0]
    ts_all = t_measured[:,1]
    te_all = t_measured[:,2]

    v1, x1, r1 = my_vel_regress(x_valid, ts_all, clean=False)
    v2, x2, r2 = my_vel_regress(x_valid, te_all, clean=False)
    if DEBUG:
        print(f'{e_time=}, v1={v1:3.1f},v2={v2:3.1f},x1={x1:4.1f},x2={x2:4.1f},r1={r1:.3f},r2={r2:.3f}')
    
    # quality control, v1,v2 大小差异，x1,x2,可信度R
    if abs(v1-v2)>20:
        continue
    if abs(x1)>800 or abs(x2)>800:
        continue
    if abs(r1)<0.9 or abs(r2)<0.9:
        continue
    
    # 方向、速度控制 
    # N2S Pass
    if N2S:
        if v1>0 or v2>0:
            continue
    if S2N:
        if v1<0 or v2<0:
            continue

    if abs(v1)>90 or abs(v2)>90:
        continue
    print('select:', e_time, abs(v1-v2),x1,x2,r1,r2)

    event_select.loc[len(event_select)]=[group_i, e_time, v1, x1, r1, v2, x2, r2]

    if DEBUG and i%100==1:
    
        traces = []
        for j,name in enumerate(HSR_STATS):
            idx_j = stats.index(name)
            traces.append(data_i[idx_j,:])
        fig, ax = plot_events(np.array(traces), x, t, fs=1, fe=20, SCALE=200/HSR_EVENTT_AMP, NORM=False)
        line1,=ax.plot(t, v1*t+x1, lw=0.5, color='r', zorder=100,label=f'x={v1:.1f}t+{x1:.1f}, R={r1:.4f}')

        line2,=ax.plot(t, v2*t+x2, lw=0.5, color='b', zorder=100,label=f'x={v2:.1f}t+{x2:.1f}, R={r2:.4f}')
        
        ax.legend(handles=[line1,line2])
        fig.tight_layout()
        fig.savefig(f'{FIG_ROOT}/{i:05d}.{e_time}.png')

csv_file = f'{H5_ROOT}/event_select.{MARKER}.csv' if not OUT_FILE else OUT_FILE

event_select.to_csv(csv_file,mode='w',index_label='idx')
print(csv_file)
    


    


        



    


